import os
import hashlib
from contextlib import contextmanager
import subprocess
from typing import Iterator, Union, Optional
import shutil

numeric = Union[int, float]


class DiskUsageSize(object):
    """
    Disk usage size

    Attributes
    ----------
    size_bytes: Union[int, float]
        size in bytes
    """

    __slots__ = "size_bytes"

    def __init__(self, size_bytes: numeric) -> None:
        if size_bytes >= 0:
            self.size_bytes: numeric = size_bytes
        else:
            raise RuntimeError("size_bytes must be positive")

    @property
    def kb(self) -> float:
        """disk usage size in kilobytes"""
        return self.size_bytes / 1024.0

    @property
    def mb(self) -> float:
        """disk usage size in megabytes"""
        return self.size_bytes / (1024 * 1024.0)

    @property
    def gb(self) -> float:
        """disk usage size in gigabytes"""
        return self.size_bytes / (1024 * 1024 * 1024.0)


def get_dir_size(dir_input: str) -> DiskUsageSize:
    """Get the size of a given directory,
    notice this could be slow if there are a lot of files under the directory.

    Parameters
    ----------
    dir_input : str
        Path to the directory to check.

    Returns
    -------
    DiskUsageSize
        Size of the directory.

    Raises
    ------
    RuntimeError
        Raised when the `dir_input` does not exist,
        or other runtime error happened in checking the size.
    """
    if not os.path.exists(dir_input):
        raise RuntimeError(f"{dir_input} does not exist")

    try:
        # -s for total size, -b to output bytes
        ret = subprocess.run(["du", "-sb", dir_input], capture_output=True)
    except Exception as err:
        raise RuntimeError(f"fail to check size for {dir_input} due to {str(err)}")

    size_bytes = int(ret.stdout.decode("utf-8").strip().split()[0])
    return DiskUsageSize(size_bytes)


def compare_dirs(dir1: str, dir2: str) -> bool:
    """Compare two directories' underlying files to make sure they are the same,
    notice this only compares content of files, ignoring file/directory metadata,
    could be slow if there are a lot of file, or there are large files under the directories.

    Parameters
    ----------
    dir1 : str
        Path to directory to be compared.
    dir2 : str
        Path to directory to be compared.

    Returns
    -------
    bool
        True if both have the same directory structure and contents, otherwise False
    """
    for root, _, file_names in os.walk(dir1, topdown=False):
        for file_name in file_names:
            file_path_dir1 = os.path.join(root, file_name)
            file_path_dir2 = os.path.join(root.replace(dir1, dir2), file_name)
            if not os.path.exists(file_path_dir2):
                return False
            else:
                md5_1 = hashlib.md5(open(file_path_dir1, "rb").read()).hexdigest()
                md5_2 = hashlib.md5(open(file_path_dir2, "rb").read()).hexdigest()
                if md5_1 != md5_2:
                    return False
    else:
        return True


def mkdir_p(dir_path: str, mode: int = 0o777, empty_exist: bool = False) -> None:
    """Make directory, following the style as in unix tool `mkdir -p`

    Parameters
    ----------
    dir_path : str
        Path of the directory to be created
    mode : int, optional
        Directory permission mode, by default 0o777
    empty_exist : bool, optional
        Empty the folder on the path if existing, by default is False

    Returns
    -------
    None
    """
    if os.path.exists(dir_path) and empty_exist:
        shutil.rmtree(dir_path)

    os.makedirs(dir_path, mode, exist_ok=True)


def in_directory(file_path: str, dir_path: str) -> bool:
    """Check if a file path is under a directory

    Parameters
    ----------
    file_path : str
        Path of the file.
    dir_path : str
        Path of the directory.

    Returns
    -------
    bool
        True if the file is under the directory, otherwise False.
    """
    dir_path = os.path.join(os.path.realpath(dir_path), "")
    file_path = os.path.realpath(file_path)
    return os.path.commonprefix([file_path, dir_path]) == dir_path


@contextmanager
def working_directory(
    dir_path: str, create_when_not_exist: bool = True
) -> Iterator[None]:
    """A context manager which changes the working directory to the given
       path, and then changes it back to its previous value on exit.

    Parameters
    ----------
    dir_path : str
        Path of the working directory.
    create_when_not_exist : bool, optional
        If True, create the directory if the given path does not exist; otherwise raise `ValueError`.
        By default True.

    Raises
    ------
    ValueError
        Raised when the given path does not exist and the option `create_when_not_exist` is False.

    Examples
    --------
    >>> # Do something in original directory
    >>> with working_directory("/my/new/path"):
    >>>     # Do something in new directory
    >>> # Back to old directory
    """
    if not os.path.exists(dir_path):
        if create_when_not_exist:
            mkdir_p(dir_path)
        else:
            raise ValueError(f"{dir_path} does not exist")

    prev_cwd = os.getcwd()
    os.chdir(dir_path)
    try:
        yield
    finally:
        os.chdir(prev_cwd)


@contextmanager
def chmod_path(
    path: str, mode: int, mode_temp: int = 0o777, child_mode: Optional[int] = None
) -> Iterator[None]:
    """A context manager which changes the permission of the given path to a temporary mode 'mode_temp'
    and then changes it to the 'mode' after processing done.

    Parameters
    ----------
    path : str
        Path of the working directory or a file.
    mode: int
        The target mode of a folder permission after processing. E.g., 0o444
    mode_temp : int, optional
        The temporary mode of a folder permission during processing. Default: 0o777
    child_mode: bool, optional
        Change permission of all child files or sub-directories under the path
        if path is a directory. Default: None not to chmod of childs.
        NOTE only the immediate child files or sub-directories will change the mode.

    Notes
    ------
    For some cases, `path` is a file, if this file does not exist, do not `chmod` its
    permission. Users may create such a file within this wrapper so that the file's
    permission will still be changed to the target `mode`.

    Examples
    --------
    >>> with chmod_path("/my/path", mode=0o555):
    >>>     # Do something in the directory
    >>>     # NOTE `path` may be a file and may not exist, after existing this wrapper
    >>>     # the file's permission will still be changed to the target `mode`.
    >>>     with open("/my/path", "w") as fout:
    >>>         # write something to the file at `/my/path`
    """
    if os.path.exists(path):
        os.chmod(path, mode_temp)

    if os.path.isdir(path) and child_mode is not None:
        for fn in os.listdir(path):
            os.chmod(os.path.join(path, fn), mode_temp)

    try:
        yield
    finally:
        if os.path.isdir(path) and child_mode is not None:
            for fn in os.listdir(path):
                os.chmod(os.path.join(path, fn), child_mode)

        if os.path.exists(path):
            os.chmod(path, mode)
